import os
from typing import Annotated

from dotenv import load_dotenv
from langchain.chat_models import init_chat_model
from langchain_tavily import TavilySearch
from langgraph.graph import StateGraph, START
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode, tools_condition
from typing_extensions import TypedDict

from com.wp.langGraph.utils import save_graph_visualization

# 加载 .env 文件
load_dotenv()
class State(TypedDict):
    messages: Annotated[list, add_messages]


llm = init_chat_model(
    model="deepseek-chat",
    temperature=0,
    model_provider="deepseek",
    api_key="sk-ea19b9ca450e46b99b38e740400ff83e",
)

graph_builder = StateGraph(State)

tool = TavilySearch(max_results=2, tavily_api_key=os.getenv('TAVILY_API_KEY'))
tools = [tool]
llm_with_tools = llm.bind_tools(tools)


def chatbot(state: State):
    return {"messages": [llm_with_tools.invoke(state["messages"])]}


graph_builder.add_node("chatbot", chatbot)

tool_node = ToolNode(tools=[tool])
graph_builder.add_node("tools", tool_node)

graph_builder.add_conditional_edges(
    "chatbot",
    tools_condition,
)
# Any time a tool is called, we return to the chatbot to decide the next step
graph_builder.add_edge("tools", "chatbot")
graph_builder.add_edge(START, "chatbot")
graph = graph_builder.compile()

# 保存状态图的可视化表示
save_graph_visualization(graph)
